import shelve
import sys
sys.path.append("/home/cyw/projects/function_sim_project/basic_script")
from easySample import easySample
import numpy as np


functionSim_predata=r"/home/cyw/projects/function_sim_project/all_data/sampleDatas/functionSim"

def get_dynamic_name(name):
    fileroot=functionSim_predata+"//"+name
    res={}
    res["adj"],res["att"],res["vtype"]=[],[],[]
    with shelve.open(fileroot) as file:
        cg=file["cg"]
        cgattr=file["cgattr"]
        caller=file["caller"]#同一行是谁调用了我
        functype=file["funcType"]
        funcs_id=file["func_id"]# name-->ind
    id_funcs={}
    for i in funcs_id.keys():
        id_funcs[funcs_id[i]]=i
    res=[]
    for i in functype.keys():
        if functype[i]=='dynamic import':
            res.append(id_funcs[i])   
    return res

def get_same_function_name(x,y,flag=True):
    """
        flag
            true  相同的
            false 不同的
    """
    name=get_dynamic_name(x)
    name1=get_dynamic_name(y)
    res=[]
    if flag:
        for i in name:
            if i in name1:
                res.append(i)
    else:
        temp=[]
        for i in name:
            if i not in name1:
                temp.append(i)
        res.append(temp)
        temp=[]
        for j in name1:
            if j not in name:
                temp.append(j)
        res.append(temp)
    return res

def get_functionSim_sample(name,filePath=r"/home/cyw/projects/function_sim_project/all_data/sampleDatas/functionSim"):
        """
            路径输入默认是ida进行反汇编
            输出：
                adj 邻接矩阵 b*n*n 
                att 特征矩阵 b*n*d
                vtype 类型矩阵 b*n*3  3种类型
        """
        fileroot=filePath+"//"+name
        res={}
        res["adj"],res["att"],res["vtype"]=[],[],[]
        with shelve.open(fileroot) as file:
            cg=file["cg"]
            cgattr=file["cgattr"]
            # caller=file["caller"]#同一行是谁调用了我
            # 这里的名称似乎写错了。。。
            try:
                functype=file["funcType"]
            except Exception as e:
                functype=file["functype"]
            try:
                funcs_id=file["func_id"]# name-->ind
            except Exception as e:
                funcs_id=file["funcs_id"]# name-->ind
        id_funcs={}
        for i in funcs_id.keys():
            id_funcs[funcs_id[i]]=i

        #   动态导入函数是[],需要处理一下,
        #   同时生成vtype
        tempData=[]
        for i in functype.keys():
            temp=[]
            if functype[i]=="local":
                temp=[1,0,0]
            elif functype[i]=='dynamic import':
                temp=[0,1,0]
                cgattr[i]=map_api_name_to_eight_bit_vector(id_funcs[i],8,20)
            else:
                temp=[0,0,1]
            tempData.append(temp)
        res["adj"]=np.array(cg)#这里是带边权的
        res["att"]=np.array(cgattr)#应该是8维，但是第一维字符串的提取似乎出现了问题，建议只是用后七维
        res["vtype"]=np.array(tempData)
        res["name"]=id_funcs
        return res

def map_api_name_to_eight_bit_vector(function_name,embeddingSize,numLim):
    """
        将api转换成一个embeddingSize位的向量，每一位使用numLim取余
    """
    function_bytes = function_name.encode('utf-8')
    eight_bit_vector = [0] * embeddingSize
    for i, byte in enumerate(function_bytes):
        index = i % embeddingSize
        eight_bit_vector[index] += byte
        eight_bit_vector[index] %= numLim
    return eight_bit_vector

def find_all_same_function(x,y):
    """
        用于判断两个样本中相同的函数有多少
        这里是粗略的判断，用嵌入值是否相同来确定
    """
    a=easySample()
    hashmap_x={}
    hashmap_y={}
    xData = get_functionSim_sample(x)
    xNames = xData["name"]
    tempData = xData["att"]
    xadj = xData["adj"]
    sample_x=[str(item)  for item in tempData]
    xSim=[0]*len(sample_x)
    for i in range(len(sample_x)):
        att = sample_x[i]
        if att not in sample_x:
            hashmap_x[att]+=i
            print("att---")
        else:
            hashmap_x[att]=i

    yData = get_functionSim_sample(y)
    yNames = yData["name"]
    tempData = yData["att"]
    yadj = yData["adj"]
    sample_y=[str(item)  for item in tempData]
    ySim=[0]*len(sample_y)
    for i in range(len(sample_y)):
        att = sample_y[i]
        if att not in sample_y:
            hashmap_y[att]+=i
            print("att++")
        else:
            hashmap_y[att]=i
    res=0
    
    for att in hashmap_x.keys():
        if att in hashmap_y:
            res+=1
            xSim[hashmap_x[att]]=1
            ySim[hashmap_y[att]]=1
            print(str(xNames[hashmap_x[att]])+"\t\t\t"+str(yNames[hashmap_y[att]]))
    print(sum(xSim),sum(ySim))
    edgeCount=0
    for i in range(len(xadj)):
        for j in range(len(xadj[i])):
            if(xadj[i][j]!=0):
                edgeCount+=1
            if(xSim[i] and xSim[j] and xadj[i][j]!=0):
                print(str(xNames[j])+"--->"+str(xNames[i]))
    print("x edgeCount={}".format(edgeCount))
    edgeCount1=0
    for i in range(len(yadj)):
        for j in range(len(yadj[i])):
            if(yadj[i][j]!=0):
                edgeCount1+=1
            if(ySim[i] and ySim[j] and yadj[i][j]!=0):
                print(str(yNames[j])+"--->"+str(yNames[i]))
    print("y edgeCount={}".format(edgeCount1))
    print("样本x长度：{} 样本y长度：{} 相同的函数个数：{}  占比：{}".format(len(sample_x),len(sample_y),res,round(res/len(sample_x),2)))
    




if __name__=="__main__":
    """
    获得两个样本中，相同的动态导入函数名称
    """
    # 异质的
    # namePair=[["d0a1ce5934f82a4ea4c82e70a9675899","69df5e279fb9c73d0efae42a72d72770"],
    #           ["df81e53f7f9555a9a45ac8a3640365c8","ffb0432940456791fc858e547a8da088"],
    #           ["77188bd267a94694d96bc9841194d03d","e3fb5232f88b4a62b0d80d31e4828770"],
    #           ["8f1901aa8ed69e2a01d169e14c11057e","c2168bda50d790c6bfa07da8571fea19"],
    #           ["9d9d2d1cb0d23a16093795152a868c42","de952d92b65848b8e2e9fb6bd37c551e"],]
    # 跨图的
    # namePair = [["f16d449b0ca61a4835577f4cef9f1624","4f5a3240fe37acab6ccb8d8469c4cb87"],
    #             ["e534dc5f410cb3677028ed2d56f20ad5","f16d449b0ca61a4835577f4cef9f1624"],
    #             ["eec55f1295d0cce40bbba7be231fdddb","6abf24e3aa3e971e6461bcfea35fe950"],
    #             ["81cbb6f4b1cc8d5e2167d106657de095","b33da0446151ac45b2e07491adde796d"],
    #             ["b68a07cf3f63172901ce8298e78d7a74","63ba29be215b532731392511dcd4363e"],
    #             ["e3c18e0bd720bd5462065988ed159955","698ff16b0e824db5874acee0f7324bb2"],
    #             ["ce1a3ebdde8eb6d24e2c91bd3ce1da72","c1eb31c7c0ff87c5b182f8e3a447028a"],
    #             ["5e9f256386b46d5d921530338af16dc0","69df5e279fb9c73d0efae42a72d72770"]]
    # 特例判断
    # namePair = [["eec55f1295d0cce40bbba7be231fdddb","6abf24e3aa3e971e6461bcfea35fe950"]]
    # for pair in namePair:
    #     x=pair[0]
    #     y=pair[1]
    #     print("******************{}***{}****************".format(x,y))
    #     res=get_same_function_name(x,y,False)
    #     print(res)
    #     print()

    # find_all_same_function("b68a07cf3f63172901ce8298e78d7a74","63ba29be215b532731392511dcd4363e")
    find_all_same_function("eec55f1295d0cce40bbba7be231fdddb","6abf24e3aa3e971e6461bcfea35fe950")
    # 74ec6006ed76a9bed14c140e53c7efd4--bb6f9c2aa88245746e24d7672016afb7  
    # find_all_same_function("74ec6006ed76a9bed14c140e53c7efd4","bb6f9c2aa88245746e24d7672016afb7")